Amazon SageMakerで複数のモデルを同じエンドポイントにホスティングしてみた

#Amazon SageMaker

#AWS

osawa yuto

2020.01.19

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

Amazon SageMakerは機械学習における開発から運用までサポートしたフルマネージドサービスです。 SageMakerでは機械学習モデルをエンドポイントにホスティングし、APIでの推論が可能です。また、複数のモデルを同じエンドポイントインスタンスにホスティングし、指定したモデルへの推論リクエストも可能です。

今回はこの複数モデルのホスティングを試してみたいと思います。

マルチモデルエンドポイント(Multi-Model Endpoints)

複数のモデルをホストできるエンドポイントを作成するには、マルチモデルエンドポイントを使用します。マルチモデルエンドポイントでは、複数のモデルをホストできる共有サービングコンテナが使用されます。これにより、単一モデルエンドポイントを使用する場合と比較して、エンドポイントの使用率が向上し、ホスティングコストが削減されます。また、Amazon SageMaker によってメモリへのモデルのロードが管理され、トラフィックパターンに基づいてモデルがスケーリングされるため、デプロイのオーバーヘッドも削減されます。マルチモデルエンドポイントは、多数のモデルをデプロイするためのスケーラブルで費用対効果の高いソリューションを提供します。

Multi-Model Endpoints で複数のモデルをホストするより

仕組み

マルチモデルエンドポイントにホスティングされたモデルの推論リクエストがあると、まずS3からモデルがエンドポイントインスタンスにダウンロードされます。ダウンロードされたモデルをメモリに読み込み、リクエストデータを推論し、推論結果をリクエスト元に返します。モデルが既にメモリに読み込まれている場合はS3からダウンロードせず、そのままメモリ上のモデルを使用します。既に読み込まれているモデルとは異なるモデルに対して推論リクエストがあり、メモリに新たにモデルを読み込む余裕がない場合は、メモリに読み込まれたモデルを解放し、必要なモデルをS3からダウンロードしてメモリに読み込みます。

次の図はモデルA、B、Cをホスティングしたマルチモデルエンドポイントの例です。モデルAに推論リクエストがあった場合はS3からモデルをダウンロードし、メモリに読み込み、推論処理を実行します。

モデルBに推論があった場合も同様です。エンドポイントインスタンスのメモリが十分にあれば、モデルAも既に読み込まれた状態です。

用途

マルチモデルエンドポイントは複数のモデルへの推論リクエストのタイミングが異なっていたり、ある程度レイテンシがあっても許容される環境やシステムに適しています。逆にレイテンシを最小限にする必要がある場合にはマルチモデルエンドポイントは向いてないため、通常通り1エンドポイントにつき1モデルをホスティングする方が適しています。

やってみる

XGBoostのモデルをマルチモデルエンドポイントでホスティングしてみます。今回はローカルのJupyter LabのPythonカーネルから実行しました。

準備

まず、バケット名やモデルを配置するS3プレフィックス、モデルのホスティングに使用するIAMロール等を定義しておきます。

import boto3
from sagemaker import get_execution_role
session = boto3.Session()
sm_client = session.client('sagemaker')
algorithm_name = 'multi-model-xgboost'

account_id = session.client('sts').get_caller_identity()['Account']
region = session.region_name
bucket = 'bucket-name'
model_prefix = 'prefix/to/models'
role = get_execution_role()

コンテナイメージの準備

XGBoostモデルをホスティングするためのマルチモデルエンドポイント用コンテナイメージを作成します。今回はSageMaker XGBoost Containerリポジトリのmmeブランチのものを利用します。

以下のスクリプトを実行して、コンテナイメージをビルドし、ECRに登録します。

%%sh -s $algorithm_name

algorithm_name=$1

account=$(aws sts get-caller-identity --query Account --output text)

# Get the region defined in the current configuration
region=$(aws configure get region)

ecr_image="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1

if [ $? -ne 0 ]
then
    aws ecr create-repository --repository-name "${algorithm_name}" > /dev/null
fi

# Get the login command from ECR and execute it directly
$(aws ecr get-login --region ${region} --no-include-email --registry-ids ${account})

# Build the docker image locally with the image name and then push it to ECR
# with the full name.

# First clear out any prior version of the cloned repo
rm -rf sagemaker-xgboost-container/

# Clone the xgboost container repo
git clone --single-branch --branch mme https://github.com/aws/sagemaker-xgboost-container.git
cd sagemaker-xgboost-container/

# Build the "base" container image that encompasses the installation of the
# XGBoost framework and all of the dependencies needed.
docker build -q -t xgboost-container-base:0.90-2-cpu-py3 -f docker/0.90-2/base/Dockerfile.cpu .

# Create the SageMaker XGBoost Container Python package.
python setup.py bdist_wheel --universal

# Build the "final" container image that encompasses the installation of the
# code that implements the SageMaker multi-model container requirements.
docker build -q -t ${algorithm_name} -f docker/0.90-2/final/Dockerfile.cpu .

docker tag ${algorithm_name} ${ecr_image}

docker push ${ecr_image}

モデルアーティファクトの準備

マルチモデルエンドポイントではホスティングするモデルを同じプレフィックスの配下に保存する必要があります。今回は事前に作成したモデルアーティファクトを同じプレフィックス配下に集めます。

※使用したモデルアーティファクトはこちらのノートブックで作成したものです。

# SageMakerの学習ジョブ名
jobs = [
    'xgb-Houston-TX-2020-01-17-03-42-01-838',
    'xgb-Chicago-IL-2020-01-17-03-41-58-963',
    'xgb-LosAngeles-CA-2020-01-17-03-41-57-386',
    'xgb-NewYork-NY-2020-01-17-03-41-56-357',
]
for job_name in jobs:
    job = sm_client.describe_training_job(TrainingJobName=job_name)
    source_uri = job['ModelArtifacts']['S3ModelArtifacts']
    target_uri = path.join('s3://', bucket, model_prefix, job['TrainingJobName'], 'model.tar.gz')
    !aws s3 cp $source_uri $target_uri

エンドポイントの作成

SageMaker Python SDKを用いることで、マルチモデルエンドポイントの作成が簡単に行えます。モデル名やモデルアーティファクトを配置したS3のプレフィックスやコンテナイメージ名、実行ロールを設定します。

import sagemaker.multidatamodel 

mdm = sagemaker.multidatamodel.MultiDataModel(
    name=model_name,
    model_data_prefix = path.join('s3://', bucket, model_prefix),
    image=f'{account_id}.dkr.ecr.{region}.amazonaws.com/{algorithm_name}:latest',
    role=role,
)
mdm.deploy(
    initial_instance_count=1,
    instance_type='ml.t2.medium'
)

インスタンス情報を設定し、デプロイします。数分ほどでデプロイが完了します。

mdm.deploy(
    initial_instance_count=1,
    instance_type='ml.t2.medium'
)

モデルの一覧を確認してみます。

models = list(mdm.list_models())
models

推論

推論しやすいように、predictorに先ほど作成したエンドポイントを設定します。

from sagemaker.predictor import csv_serializer
predictor = sagemaker.predictor.RealTimePredictor(
    endpoint=mdm.endpoint_name,
    serializer=csv_serializer
)

推論するデータと対象モデルを指定して推論してみます。

predictor.predict(
    feature,
    target_model=models[0]
)

今回のケースだと、初回の推論リクエストのレイテンシは1秒程度、２回目以降だと数十ミリ秒から数百ミリ秒程度でした。モデルアーティファクトのサイズや処理内容に応じてこの時間は大きく変わりそうです。

片付け

余分な費用がかからないように、使用し終わったらエンドポイントインスタンスを削除します。

predictor.delete_endpoint()

さいごに

Amazon SageMakerで複数のモデルを1つのエンドポイントでホスティングするマルチモデルエンドポイントについて紹介しました。開発環境やレイテンシが許容できるシステム等においてはコストを節約できる良い方法だと思います。ただし、通常の推論エンドポイント用のコンテナイメージそのままでは使用できないため、コンテナイメージや周辺システムの修正コストとマルチモデルエンドポイントによって節約できるコストのバランスを見ながら、利用を検討すると良さそうです。